from clickhouse_driver import Client
import pandas as pd

CH = Client(host='clickhouse', port=9000, database='qbot')

def rps(period: int = 5) -> pd.DataFrame:
    """返回周期 RPS 排名 0-100"""
    df = CH.query_dataframe(f'''
        SELECT ind,
               (close - first_value(close) OVER (PARTITION BY ind ORDER BY dt
                ROWS BETWEEN {period * 240} PRECEDING AND CURRENT ROW))
               / first_value(close) OVER (PARTITION BY ind ORDER BY dt
                ROWS BETWEEN {period * 240} PRECEDING AND CURRENT ROW) AS ret
        FROM sector_1m
        WHERE dt >= now() - INTERVAL {period} DAY
    ''')
    df['rps'] = (1 - df.groupby('ret')['ret'].rank(pct=True)) * 100
    return df.groupby('ind')['rps'].last().reset_index().sort_values('rps', ascending=False)

def money_flow(ind: str) -> float:
    """当日主力净流入（万元）"""
    df = ak.stock_board_industry_cons_em(symbol=ind)
    if df.empty:
        return 0
    # 用板块成分股主力净流入求和
    return ak.stock_individual_fund_flow(stock=df['代码'].tolist(), indicator='主力净流入').sum()
